import torch
import torch.nn.functional as F
import os
import argparse
from tqdm import trange
from transformers import GPT2LMHeadModel, GPT2Config, CpmTokenizer
from utils import top_k_top_p_filtering, set_logger
from gpt_config import GPTConfig,gpt_logger
from os.path import join, exists
import ai
args = GPTConfig()
device = 'cuda:0' if args.cuda else 'cpu'
# device = 'cpu'

# 创建日志对象
logger = gpt_logger

# 初始化tokenizer
tokenizer = CpmTokenizer(vocab_file=args.vocab_file_path)
eod_id = tokenizer.convert_tokens_to_ids("<eod>")  # 文档结束符
sep_id = tokenizer.sep_token_id
unk_id = tokenizer.unk_token_id

# 加载模型
model = GPT2LMHeadModel.from_pretrained(args.pretrained_model)
model.eval()
model = model.to(device)


def generate_next_token(input_ids):
    """
    对于给定的上文，生成下一个单词
    """
    outputs = model(input_ids=input_ids)
    logits = outputs.logits
    # next_token_logits表示最后一个token的hidden_state对应的prediction_scores,也就是模型要预测的下一个token的概率
    next_token_logits = logits[0, -1, :]
    next_token_logits = next_token_logits / args.temperature
    # 对于<unk>的概率设为无穷小，也就是说模型的预测结果不可能是[UNK]这个token
    next_token_logits[unk_id] = -float('Inf')
    filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=args.topk, top_p=args.topp)
    # torch.multinomial表示从候选集合中选出无放回地进行抽取num_samples个元素，权重越高，抽到的几率越高，返回元素的下标
    next_token_id = torch.multinomial(F.softmax(filtered_logits, dim=-1), num_samples=1)
    return next_token_id


def generate(max_len):
    # 对title与context进行tokenize
    title_ids = tokenizer.encode(title, add_special_tokens=False)
    context_ids = tokenizer.encode(context, add_special_tokens=False)
    input_ids = title_ids + [sep_id] + context_ids
    cur_len = len(input_ids)
    last_token_id = input_ids[-1]  # 已生成的内容的最后一个token
    input_ids = torch.tensor([input_ids], dtype=torch.long, device=device)

    while True:
        next_token_id = generate_next_token(input_ids[:, -args.context_len:])
        input_ids = torch.cat((input_ids, next_token_id.unsqueeze(0)), dim=1)
        cur_len += 1
        word = tokenizer.convert_ids_to_tokens(next_token_id.item())
        # if cur_len >= max_len:
        #     break
        # 超过最大长度，并且换行
        if cur_len >= max_len and last_token_id == 8 and next_token_id == 3:
            break
        # 超过最大长度，并且生成标点符号
        if cur_len >= max_len and word in [".", "。", "！", "!", "?", "？", ",", "，"]:
            break
        # 生成结束符
        if next_token_id == eod_id:
            break
    result = tokenizer.decode(input_ids.squeeze(0))
    return result


if __name__ == '__main__':
    title = "我的家乡真美"
    context = "每个角落都流淌着我童年的记忆。在我家的小院里，有一棵高大的榕树，它的树冠像一把巨大的绿色伞 "
    # logger.info("title:{}".format(title))
    # logger.info("context:{}".format(context))

    logger.info('开始生成')
    # print(model)

    # 开始生成
    result = generate(100)
    result = result.split("<sep>")[1]
    logger.info("result:{}\n".format(result))
